#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
HHL (integrated) — uses the user's own modules directly
=======================================================

Modules used exactly as provided:
- quantum_state_preparation.load_vector_from_qRAM
- Hamiltonian_simulator.build_pauli_evolution_oracle_from_matrix
- qpe_hhl.hhl_pe_encode / hhl_pe_uncompute / pick_t_from_lambda_max / rough_lambda_max_via_eigs
- lookup_rotation.LookupRotation

This file exposes:
- build_hhl_circuit_user(A, b, t_bits, method='suzuki2', num_slices=2, lambda_min=None, negative_evals=False)
- helper: postselect_statevector_on_flag(sv, total_qubits, flag_index)

And a demo main() that prints stepwise info when you run:
    python hhl_user_integrated.py
"""
from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Tuple, Optional

import numpy as np
from numpy.linalg import norm

from qiskit import QuantumCircuit, QuantumRegister
from qiskit.quantum_info import Statevector

# --- import your modules ---
from algorithm.qiskit.quantum_state_preparation import load_vector_from_qRAM
from algorithm.qiskit.Hamiltonian_simulator import (
    build_pauli_evolution_oracle_from_matrix,
    expm_hermitian,
    block_diag_control_U,
    unitary_global_distance,
)
from algorithm.qiskit.qpe_hhl import (
    hhl_pe_encode, hhl_pe_uncompute,
    pick_t_from_lambda_max, rough_lambda_max_via_eigs
)
from algorithm.qiskit.lookup_rotation import LookupRotation


# =========================
# HHL (user modules integrated) — complete pipeline + post-processing
# =========================
from dataclasses import dataclass
from typing import Optional, Tuple, Dict, Any, List
import numpy as np
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.quantum_info import Statevector

# ---------- 若你未粘贴 _append_by_qreg，可用这个版本 ----------
def _append_by_qreg(target: QuantumCircuit, other: QuantumCircuit) -> None:
    """将 other 按量子/经典寄存器对齐追加到 target（兼容多寄存器电路拼接）。"""
    for qreg in other.qregs:
        if qreg not in target.qregs:
            target.add_register(qreg)
    for creg in other.cregs:
        if creg not in target.cregs:
            target.add_register(creg)
    data_copy = other.data.copy() if other is target else other.data
    for inst in data_copy:
        target._append(*inst)
    target.global_phase += other.global_phase
# -----------------------------------------------------------

# ---------- 工具：维度扩展/厄米化 与 尺寸回退 ----------
def expand_to_hermitian(A: np.ndarray, b: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """把一般（可能非厄米）方阵 A 扩展为厄米块矩阵 H，并相应扩展 b.
       H = [[0, A^H],
            [A, 0]]
       b' = [0, b]
    """
    A = np.asarray(A, dtype=complex)
    b = np.asarray(b, dtype=complex).reshape(-1)
    m = A.shape[0]
    Z = np.zeros_like(A, dtype=complex)
    H = np.block([[Z, A.conj().T],
                  [A, Z]])
    b2 = np.concatenate([np.zeros(m, dtype=complex), b])
    return H, b2

def expand_to_powerdim(matrix: np.ndarray, vector: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """扩展到最近的 2^n 维：矩阵在对角补 1，向量在尾部补 0。"""
    mat_dim = matrix.shape[0]
    n = int(np.ceil(np.log2(mat_dim)))
    newN = 2 ** n
    newM = np.eye(newN, dtype=complex)
    newM[:mat_dim, :mat_dim] = matrix
    newV = np.zeros(newN, dtype=complex)
    newV[:len(vector)] = vector
    return newM, newV

def matrix_resize_if_needed(A: np.ndarray, b: np.ndarray) -> Tuple[np.ndarray, np.ndarray, bool, bool, int]:
    """返回 (A2, b2, truncated_powerdim, truncated_hermitian, orig_size)"""
    # ---- 新增：确保 A,b 是 ndarray，且 b 为 1D 向量 ----
    if not isinstance(A, np.ndarray):
        A = np.asarray(A)
    if not isinstance(b, np.ndarray):
        b = np.asarray(b)
    b = b.reshape(-1).astype(complex)   # <<< 关键：把 (1,2)/(2,1) 等统一成 (2,)

    if A.shape[0] != A.shape[1]:
        raise ValueError("输入矩阵必须是方阵")
    if A.shape[0] != len(b):
        raise ValueError("矩阵与向量维度不一致")

    orig_size = len(b)
    truncated_powerdim = False
    truncated_hermitian = False

    # 维度到 2^n
    is_power = (np.log2(A.shape[0]) % 1) == 0
    if not is_power:
        A, b = expand_to_powerdim(A, b)
        truncated_powerdim = True

    # 厄米性
    is_herm = np.allclose(A, A.conj().T)
    if not is_herm:
        A, b = expand_to_hermitian(A, b)
        truncated_hermitian = True

    return A, b, truncated_powerdim, truncated_hermitian, orig_size


def _resize_back_vector(vec: np.ndarray, truncated_powerdim: bool, truncated_hermitian: bool, orig_size: int) -> np.ndarray:
    """将结果向量裁剪回原问题维度。"""
    out = vec.copy()
    if truncated_hermitian:
        # [[0, A^H],[A,0]] 的解向量前半部分是扩展引入的（与 b* 对应），取后半
        half = out.shape[0] // 2
        out = out[half:]
    if truncated_powerdim:
        out = out[:orig_size]
    return out

def _resize_back_matrix(M: np.ndarray, truncated_powerdim: bool, truncated_hermitian: bool, orig_size: int) -> np.ndarray:
    """将矩阵裁剪回原问题维度。"""
    A = M.copy()
    if truncated_hermitian:
        # 从 [[0,A^H],[A,0]] 还原 A
        full = A.shape[0]; half = full // 2
        A = A[half:, :half]
    if truncated_powerdim:
        A = A[:orig_size, :orig_size]
    return A

# ---------- 工具：从全态向量中抽取 anc=1 & QPE寄存器归零 后的工作寄存器态 ----------
def _extract_solution_state(qc, sv, work, phase, rec):
    # 以 qc.qubits 的顺序作为张量轴顺序
    qubits = list(qc.qubits)
    axis_of = {qb: qubits.index(qb) for qb in qubits}
    N = len(qubits)

    tensor = sv.data.reshape([2]*N)

    # 固定 anc=1，phase=0...0，lookup 的 work/msq=0
    fixed = {
        axis_of[rec._anc[0]]: 1,
        axis_of[rec._workq[0]]: 0,
        axis_of[rec._msq[0]]: 0,
    }
    for qb in phase:
        fixed[axis_of[qb]] = 0

    # 其余轴自由（工作寄存器）
    idx = [slice(None)]*N
    for ax, val in fixed.items():
        idx[ax] = val

    sub = tensor[tuple(idx)]
    out = sub.reshape(-1).astype(complex)

    # 期望长度=2**len(work)
    expect_len = 2 ** len(work)
    if out.shape[0] != expect_len:  # 强力兜底：如果轴对齐有误，直接报错，避免静默出错
        raise RuntimeError(
            f"Extraction length {out.shape[0]} != expected {expect_len}. "
            "Check register wiring and projection axes."
        )

    p_succ = float(np.vdot(out, out).real)
    return out, p_succ

# ---------- 结果结构 ----------
@dataclass
class HHLUserResult:
    solution: np.ndarray             # 重标定后的解向量（回到原问题维度）
    output_unscaled: np.ndarray      # 投影态（工作寄存器），归一化前
    probability_success: float       # anc=1 & registers=0 的概率
    matrix_resized: np.ndarray       # （回到原维度后的）矩阵
    vector_resized: np.ndarray       # （回到原维度后的）右端向量
    circuit: QuantumCircuit          # 完整电路

# ---------- 主构建函数 ----------
def build_hhl_circuit_user(
    A: np.ndarray,
    b: np.ndarray,
    t_bits: int = 6,
    method: str = 'suzuki2',
    num_slices: int = 2,
    negative_evals: bool = False,
    lambda_min: Optional[float] = None,
    t: Optional[float] = None,
    scale: float = 0.0,  # 可留 0 交由 LookupRotation 内部根据 (lambda_min, t) 推断
    measure_success: bool = False
) -> Tuple[QuantumCircuit, Dict[str, Any]]:
    """
    组装 HHL 主电路，不做执行。返回 (qc, meta)，其中 meta 含寄存器引用便于后处理。
    """
    # 1) 维度/厄米处理
    A2, b2, trunc_pow, trunc_herm, orig_size = matrix_resize_if_needed(A, b)
    N = A2.shape[0]
    n = int(np.round(np.log2(N)))
    if 2 ** n != N:
        raise ValueError("内部错误：扩展后维度非 2^n")

    # 2) 选择演化时间 t
    if t is None:
        lam_max = rough_lambda_max_via_eigs(A2)
        t = pick_t_from_lambda_max(lam_max)

    # 3) 构建受控演化 oracle（内含恒等项相位）
    evolution_oracle = build_pauli_evolution_oracle_from_matrix(A2, method=method, num_slices=num_slices)

    # 4) 电路与寄存器
    work = QuantumRegister(n, "wrk")           # A 的工作寄存器
    phase = QuantumRegister(t_bits, "ph")      # 相位估计寄存器
    qc = QuantumCircuit(work, phase)

    # 5) |b> 初态制备（用你的 qRAM 加载）
    load_vector_from_qRAM(list(b2.astype(complex)), qc=qc, pos_init=0)

    # 6) 相位估计编码（无测量）
    hhl_pe_encode(qc, phase, work, evolution_oracle, t)

    # 7) 查表倒数旋转（LookupRotation）
    rec = LookupRotation(
        pat_length=None, subpat_length=None,
        scale=scale, negative_evals=negative_evals,
        evo_time=t, lambda_min=lambda_min
    )
    rec_circ = rec.construct_circuit("circuit", phase[::-1])  # 绑定相位寄存器
    _append_by_qreg(qc, rec_circ)

    # 8) 逆相位估计（删除相位寄存器中的信息）
    hhl_pe_uncompute(qc, phase, work, evolution_oracle, t)

    # 9) 可选：测成功位（anc）
    if measure_success:
        c = ClassicalRegister(1, "c_succ")
        qc.add_register(c)
        qc.measure(rec._anc, c)

    meta = dict(
    A2=A2, b2=b2,
    trunc_pow=trunc_pow, trunc_herm=trunc_herm, orig_size=orig_size,
    work=work, phase=phase, rec=rec, t=t, method=method, num_slices=num_slices
    )
    return qc, meta

# ---------- 运行（Statevector）并做后处理 ----------
def run_hhl_user_statevector(
    A: np.ndarray,
    b: np.ndarray,
    t_bits: int = 6,
    method: str = 'suzuki2',
    num_slices: int = 2,
    negative_evals: bool = False,
    lambda_min: Optional[float] = None,
    t: Optional[float] = None,
    scale: float = 0.0
) -> HHLUserResult:
    """
    使用 Statevector 精确仿真（无噪声）得到 HHL 解及后处理。
    """
    qc, meta = build_hhl_circuit_user(
        A, b, t_bits=t_bits, method=method, num_slices=num_slices,
        negative_evals=negative_evals, lambda_min=lambda_min, t=t, scale=scale,
        measure_success=False
    )

    # 1) 精确态矢
    sv = Statevector.from_instruction(qc)

    # 2) 抽取 anc=1 & 其他寄存器=0 的工作寄存器态，得到成功概率
    out_unscaled, p_succ = _extract_solution_state(
        qc, sv, work=meta["work"], phase=meta["phase"], rec=meta["rec"]
    )

    # 3) 归一化输出（工作寄存器）
    if p_succ > 0:
        out_normed = out_unscaled / np.sqrt(p_succ)
    else:
        out_normed = out_unscaled.copy()

    # 4) 回到原维度后的矩阵/向量
    A_resized = _resize_back_matrix(meta["A2"], meta["trunc_pow"], meta["trunc_herm"], meta["orig_size"])
    b_resized = np.asarray(b).reshape(-1).astype(complex)

    # 将 out_normed 同样回到原维度（裁剪前需考虑：work 的逻辑维度= A2 的维度）
    x_hat_full = _resize_back_vector(out_normed, meta["trunc_pow"], meta["trunc_herm"], meta["orig_size"])

    # 5) HHL 常用重标定（与 Qiskit Aqua 相同的做法）：
    #    令 tmp = A x_hat，f1 = ||b|| / ||tmp||，f2 校准整体相位，然后 x = f1 * e^{-i f2} x_hat
    tmp = A_resized.dot(x_hat_full)
    if np.linalg.norm(tmp) > 0:
        f1 = np.linalg.norm(b_resized) / np.linalg.norm(tmp)
    else:
        f1 = 1.0
    # 避免 -0.-0.j 的相位毛刺，参考 Aqua 里 -1+1 的 trick
    with np.errstate(invalid='ignore'):
        phi = np.angle(b_resized * np.conj(tmp) - 1 + 1)
    f2 = float(np.sum(phi) / max(1, int(np.log2(A_resized.shape[0]))))
    solution = f1 * x_hat_full * np.exp(-1j * f2)


    return HHLUserResult(
        solution=solution,
        output_unscaled=out_unscaled,
        probability_success=p_succ,
        matrix_resized=A_resized,
        vector_resized=b_resized,
        circuit=qc
    )

# ---------- Demo ----------
if __name__ == "__main__":
    # 一个小示例：2x2 SPD
    A_demo = np.array([[1, 0], [0, 1]], dtype=complex)
    b_demo = np.array([1.0, 1.0], dtype=complex)

    res = run_hhl_user_statevector(
        A_demo, b_demo,
        t_bits=3,
        method='suzuki2',
        num_slices=16,
        negative_evals=False,
        lambda_min=None,   # 也可给出最小特征值以便 LookupRotation 内部设定 scale
        t=None,            # 自动根据 λ_max 选 t
        scale=0.0          # 0 表示由 LookupRotation 自动推断
    )

    print("=== HHL (user modules) Result ===")
    print(f"Success prob = {res.probability_success:.6f}")
    print("Solution (rescaled to original dim):")
    print(res.solution)
    # 验证 A x ≈ b
    print("A @ x:", res.matrix_resized @ res.solution)
    print("b    :", res.vector_resized)
    print("shapes:", res.matrix_resized.shape, res.solution.shape)
    print("A@x (ravel):", (res.matrix_resized @ res.solution).ravel())
    print("||Ax - b||:", np.linalg.norm(res.matrix_resized @ res.solution - res.vector_resized))

# import numpy as np
# from numpy.linalg import eigvalsh, norm
# import pytest


# @pytest.mark.parametrize("A,b", [
#     # 简单 2x2 SPD：易于相位估计聚相
#     (np.array([[1.5, 0.5], [0.5, 1.5]], dtype=complex),
#      np.array([1.0, 0.0], dtype=complex)),

#     # 4x4 SPD：稍大一点，覆盖多比特 work 与更多相位位数
#     (np.array([[4.0, 1.0, 0.0, 0.0],
#                [1.0, 3.5, 0.2, 0.0],
#                [0.0, 0.2, 2.5, 0.7],
#                [0.0, 0.0, 0.7, 2.2]], dtype=complex),
#      np.array([1.0, 0.0, 0.0, 0.0], dtype=complex)),
# ])
# def test_hhl_solve_linear_system(A, b):
#     # 为了提高 LookupRotation 的稳定性，给出 lambda_min
#     lam = eigvalsh((A + A.conj().T)/2.0)  # 保守取 Hermitian 部分
#     lambda_min = float(np.min(np.abs(lam)))
#     assert lambda_min > 1e-9, "Matrix is (near-)singular for HHL."

#     # 跑 HHL（无噪声 statevector）
#     res = run_hhl_user_statevector(
#         A, b,
#         t_bits=8,           # 相位位数
#         method='suzuki2',   # 二阶 Suzuki
#         num_slices=16,
#         negative_evals=False,
#         lambda_min=lambda_min,
#         t=None,             # 由谱半径自动选择
#         scale=0.0           # 让 LookupRotation 根据 (lambda_min, t) 自定缩放
#     )

#     x_hhl = res.solution
#     x_true = np.linalg.solve(A, b)

#     # 对齐整体相位再比较
#     phase = np.angle(np.vdot(x_true, x_hhl))
#     x_aligned = x_hhl * np.exp(-1j * phase)

#     rel_err = norm(x_aligned - x_true) / max(1e-12, norm(x_true))
#     resid   = norm(A.dot(x_aligned) - b) / max(1e-12, norm(b))

#     # 经验阈值（取决于 t_bits、Trotter slice 等）
#     assert rel_err < 5e-2, f"Relative error too large: {rel_err:.3e}"
#     assert resid   < 5e-2, f"Residual too large: {resid:.3e}"

#     # 成功率要合理（不是零）
#     assert res.probability_success > 1e-6